/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2020. All rights reserved.
 * Licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 * Description: C file template for CA
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include "tee_client_api.h"
#include "sign.h"
#include "randombytes.h"
#include <stdint.h>
#include <time.h>

#define VERSION_BUFFER_SIZE 256
#define OPERATION_START_FLAG 1
#define OUT_BUFF_INDEX 3

enum {
	DILITHIUM_NIST2,
	DILITHIUM_NIST3,
	DILITHIUM_NIST5
};

// #define DILITHIUM2_CRYPTO_PUBLICKEYBYTES 1312
// #define DILITHIUM2_CRYPTO_SECRETKEYBYTES 2544
// #define DILITHIUM2_CRYPTO_BYTES 2420

// #define DILITHIUM3_CRYPTO_PUBLICKEYBYTES 1952
// #define DILITHIUM3_CRYPTO_SECRETKEYBYTES 4016
// #define DILITHIUM3_CRYPTO_BYTES 3293

// #define DILITHIUM5_CRYPTO_PUBLICKEYBYTES 2592
// #define DILITHIUM5_CRYPTO_SECRETKEYBYTES 4880
// #define DILITHIUM5_CRYPTO_BYTES 4595

#define DILITHIUM_MODE 2

#if DILITHIUM_MODE == 2
#define CRYPTO_PUBLICKEYBYTES 1312
#define CRYPTO_SECRETKEYBYTES 2544
#define CRYPTO_BYTES 2420

#elif DILITHIUM_MODE == 3
#define CRYPTO_PUBLICKEYBYTES 1952
#define CRYPTO_SECRETKEYBYTES 4016
#define CRYPTO_BYTES 3293

#elif DILITHIUM_MODE == 5
#define CRYPTO_PUBLICKEYBYTES 2592
#define CRYPTO_SECRETKEYBYTES 4880
#define CRYPTO_BYTES 4595
#endif

static const TEEC_UUID g_demoTemplateUuid = {
    0xad9ceb1f, 0xc3f6, 0x4de7, { 0xa1, 0x2a, 0x9a, 0xeb, 0x79, 0x7c, 0xec, 0x10 }
};

static uint8_t *k_qta_path = (uint8_t*)
    "/home/pqc-demo/ad9ceb1f-c3f6-4de7-a12a-9aeb797cec10.sec";

enum {
    CMD_GET_TA_VERSION = 1,
    CMD_DILITHIUM_KEYGEN = 2,
    CMD_DILITHIUM_SIGN = 3,
    CMD_DILITHIUM_VERIFY = 4,
    CMD_DILITHIUM_INIT = 5,
    CMD_DILITHIUM_TESTBENCH = 6,
};

int get_ta_version()
{
    TEEC_Context context = {0};
    TEEC_Session session = {0};
    TEEC_Result result;
    TEEC_Operation operation = {0};
    uint32_t origin = 0;

    char versionBuf[VERSION_BUFFER_SIZE] = {0};
    unsigned int bufLen = VERSION_BUFFER_SIZE;
    result = TEEC_InitializeContext(NULL, &context);
    if (result != TEEC_SUCCESS) {
        printf("teec initial failed");
        goto cleanup_1;
    }


    context.ta_path = k_qta_path;

    /* MUST use TEEC_LOGIN_IDENTIFY method */
    operation.started = 1;
    operation.paramTypes = TEEC_PARAM_TYPES(
        TEEC_NONE,
        TEEC_NONE,
        TEEC_NONE,
        TEEC_NONE);

    result = TEEC_OpenSession(
        &context, &session, &g_demoTemplateUuid, TEEC_LOGIN_IDENTIFY, NULL, &operation, NULL);
    if (result != TEEC_SUCCESS) {
        printf("\nteec open session failed\n");
        goto cleanup_2;
    }

    operation.started = OPERATION_START_FLAG;
    operation.paramTypes = TEEC_PARAM_TYPES(
        TEEC_NONE,
        TEEC_NONE,
        TEEC_NONE,
        TEEC_MEMREF_TEMP_OUTPUT);
    operation.params[OUT_BUFF_INDEX].tmpref.buffer = versionBuf;
    operation.params[OUT_BUFF_INDEX].tmpref.size = bufLen;

    result = TEEC_InvokeCommand(&session, CMD_GET_TA_VERSION, &operation, &origin);
    if (result != TEEC_SUCCESS) {
        printf("invoke failed, codes=0x%x, origin=0x%x", result, origin);
    } else {
        printf("Succeed to load TA, TA's version: %s.\n", versionBuf);
    }

    TEEC_CloseSession(&session);
cleanup_2:
    TEEC_FinalizeContext(&context);
cleanup_1:
    return 0;
}

TEEC_Result init_dilithium_level(int level)
{
    printf("***************** init_dilithium_level ********\n");
    TEEC_Context context = {0};
    TEEC_Session session = {0};
    TEEC_Result result;
    TEEC_Operation operation = {0};
    uint32_t origin = 0;

    TEEC_SharedMemory levelSM;

    /* ========================================================================
    [1] Connect to TEE
    ======================================================================== */
    result = TEEC_InitializeContext(NULL, &context);
    if (result != TEEC_SUCCESS)
    {
        printf("teec initial failed\n");
        goto cleanup_1;
    }

    /* ========================================================================
    [2] Open session with TEE application
    ======================================================================== */
    /* Open a Session with the TEE application. */
    /* MUST use TEEC_LOGIN_IDENTIFY method */
    operation.started = OPERATION_START_FLAG;
    operation.paramTypes = TEEC_PARAM_TYPES(TEEC_NONE, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    context.ta_path = k_qta_path;
    result = TEEC_OpenSession(&context, &session, &g_demoTemplateUuid, TEEC_LOGIN_IDENTIFY, NULL, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("teec open session failed\n");
        goto cleanup_2;
    }

    /* ========================================================================
    [3] Initialize the Shared Memory buffers
    ======================================================================== */
    /* Bulk input buffer. */
    levelSM.size = sizeof(int);
    levelSM.flags = TEEC_MEM_INPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    levelSM.buffer = (uint8_t*)(&level);
    result = TEEC_RegisterSharedMemory(&context, &levelSM);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_3;
    }

    /* Start the dilithium sign operation within the TEE application. */
    operation.paramTypes =
        TEEC_PARAM_TYPES(TEEC_MEMREF_PARTIAL_INPUT, TEEC_NONE, TEEC_NONE, TEEC_NONE);

    operation.params[0].memref.parent = &levelSM;
    operation.params[0].memref.offset = 0;
    operation.params[0].memref.size = sizeof(int);

    /* ========================================================================
    [4] Perform cryptographic operation dilithium signature
    ======================================================================== */
    result = TEEC_InvokeCommand(&session, CMD_DILITHIUM_INIT, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("Initializing the dilithium level failed\n");
        goto cleanup_4;
    }
    else
    {
        printf("Initializing the dilithium level succeeded\n");
    }

    TEEC_CloseSession(&session);
    return TEEC_SUCCESS;

cleanup_4:
    TEEC_ReleaseSharedMemory(&levelSM);
cleanup_3:
    TEEC_CloseSession(&session);
cleanup_2:
    TEEC_FinalizeContext(&context);
    return result;
cleanup_1:
    return result;
}


TEEC_Result get_dilithium_keys(uint8_t *pk, uint32_t pkLen, uint8_t *sk,
                               uint32_t skLen)
{
    TEEC_Context context = {0};
    TEEC_Session session = {0};
    TEEC_Result result;
    TEEC_Operation operation = {0};
    uint32_t origin = 0;

    TEEC_SharedMemory pkSM;
    TEEC_SharedMemory skSM;
    /* ========================================================================
    [1] Connect to TEE
    ======================================================================== */
    result = TEEC_InitializeContext(NULL, &context);
    if (result != TEEC_SUCCESS)
    {
        printf("teec initial failed\n");
        goto cleanup_1;
    }

    /* ========================================================================
    [2] Open session with TEE application
    ======================================================================== */
    /* Open a Session with the TEE application. */
    /* MUST use TEEC_LOGIN_IDENTIFY method */
    operation.started = OPERATION_START_FLAG;
    operation.paramTypes = TEEC_PARAM_TYPES(TEEC_NONE, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    context.ta_path = k_qta_path;
    result = TEEC_OpenSession(&context, &session, &g_demoTemplateUuid, TEEC_LOGIN_IDENTIFY, NULL, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("teec open session failed\n");
        goto cleanup_2;
    }

    /* ========================================================================
    [3] Initialize the Shared Memory buffers
    ======================================================================== */
    /* Bulk input buffer. */
    pkSM.size = pkLen;
    pkSM.flags = TEEC_MEM_OUTPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    pkSM.buffer = pk;
    result = TEEC_RegisterSharedMemory(&context, &pkSM);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_3;
    }
    /* Bulk output buffer. */
    skSM.size = skLen;
    skSM.flags = TEEC_MEM_OUTPUT;
    skSM.buffer = sk;
    /* Use TEE Client API to register the underlying memory buffer. */
    result = TEEC_RegisterSharedMemory(&context, &skSM);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_4;
    }

    /* Start the dilithium sign operation within the TEE application. */
    operation.paramTypes =
        TEEC_PARAM_TYPES(TEEC_MEMREF_PARTIAL_OUTPUT, TEEC_MEMREF_PARTIAL_OUTPUT, TEEC_NONE, TEEC_NONE);

    operation.params[0].memref.parent = &pkSM;
    operation.params[0].memref.offset = 0;
    operation.params[0].memref.size = pkLen;

    operation.params[1].memref.parent = &skSM;
    operation.params[1].memref.offset = 0;
    operation.params[1].memref.size = skLen;

    /* ========================================================================
    [4] Perform cryptographic operation dilithium signature
    ======================================================================== */
    result = TEEC_InvokeCommand(&session, CMD_DILITHIUM_KEYGEN, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("key generation using dilithium failed\n");
        goto cleanup_5;
    }
    else
    {
        printf("key generation using dilithium succeeded\n");
    }

    TEEC_CloseSession(&session);
    return TEEC_SUCCESS;

cleanup_5:
    TEEC_ReleaseSharedMemory(&skSM);
cleanup_4:
    TEEC_ReleaseSharedMemory(&pkSM);
cleanup_3:
    TEEC_CloseSession(&session);
cleanup_2:
    TEEC_FinalizeContext(&context);
    return result;
cleanup_1:
    return result;
}

TEEC_Result get_dilithium_signature(uint8_t const *msg, uint32_t msgLen, uint8_t const *sk, uint32_t skLen, uint8_t *sig,
                                    uint32_t sigLen)
{
    TEEC_Context context = {0};
    TEEC_Session session = {0};
    TEEC_Result result;
    TEEC_Operation operation = {0};
    uint32_t origin = 0;

    TEEC_SharedMemory msgSm;
    TEEC_SharedMemory skSm;
    TEEC_SharedMemory sigSm;
    /* ========================================================================
    [1] Connect to TEE
    ======================================================================== */
    result = TEEC_InitializeContext(NULL, &context);
    if (result != TEEC_SUCCESS)
    {
        printf("teec initial failed\n");
        goto cleanup_1;
    }

    /* ========================================================================
    [2] Open session with TEE application
    ======================================================================== */
    /* Open a Session with the TEE application. */
    /* MUST use TEEC_LOGIN_IDENTIFY method */
    operation.started = OPERATION_START_FLAG;
    operation.paramTypes = TEEC_PARAM_TYPES(TEEC_NONE, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    context.ta_path = k_qta_path;
    result = TEEC_OpenSession(&context, &session, &g_demoTemplateUuid, TEEC_LOGIN_IDENTIFY, NULL, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("teec open session failed\n");
        goto cleanup_2;
    }

    /* ========================================================================
    [3] Initialize the Shared Memory buffers
    ======================================================================== */
    /* Bulk input buffer. */
    msgSm.size = msgLen;
    msgSm.flags = TEEC_MEM_INPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    msgSm.buffer = (uint8_t *)msg;
    result = TEEC_RegisterSharedMemory(&context, &msgSm);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_3;
    }

    skSm.size = skLen;
    skSm.flags = TEEC_MEM_INPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    skSm.buffer = (uint8_t *)sk;
    result = TEEC_RegisterSharedMemory(&context, &skSm);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_4;
    }
    /* Bulk output buffer. */
    sigSm.size = sigLen;
    sigSm.flags = TEEC_MEM_OUTPUT;
    sigSm.buffer = sig;
    /* Use TEE Client API to register the underlying memory buffer. */
    result = TEEC_RegisterSharedMemory(&context, &sigSm);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_5;
    }

    /* Start the dilithium sign operation within the TEE application. */
    operation.paramTypes =
        TEEC_PARAM_TYPES(TEEC_MEMREF_PARTIAL_INPUT, TEEC_MEMREF_PARTIAL_INPUT, TEEC_MEMREF_PARTIAL_OUTPUT, TEEC_NONE);

    operation.params[0].memref.parent = &msgSm;
    operation.params[0].memref.offset = 0;
    operation.params[0].memref.size = msgLen;

    operation.params[1].memref.parent = &skSm;
    operation.params[1].memref.offset = 0;
    operation.params[1].memref.size = skLen;

    operation.params[2].memref.parent = &sigSm;
    operation.params[2].memref.offset = 0;
    operation.params[2].memref.size = sigLen;

    /* ========================================================================
    [4] Perform cryptographic operation dilithium signature
    ======================================================================== */
    result = TEEC_InvokeCommand(&session, CMD_DILITHIUM_SIGN, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("digital signature using dilithium failed\n");
        goto cleanup_6;
    }
    else
    {
        printf("digital signature using dilithium succeeded\n");
    }

    TEEC_CloseSession(&session);
    return TEEC_SUCCESS;

cleanup_6:
    TEEC_ReleaseSharedMemory(&sigSm);
    return result;
cleanup_5:
    TEEC_ReleaseSharedMemory(&skSm);
    return result;
cleanup_4:
    TEEC_ReleaseSharedMemory(&msgSm);
    return result;
cleanup_3:
    TEEC_CloseSession(&session);
    return result;
cleanup_2:
    TEEC_FinalizeContext(&context);
    return result;
cleanup_1:
    return result;
}


TEEC_Result get_dilithium_verify_result(uint8_t const *pk, uint32_t pkLen, uint8_t *sm, uint32_t smLen, int *verify_result_ptr)
{
    TEEC_Context context = {0};
    TEEC_Session session = {0};
    TEEC_Result result;
    TEEC_Operation operation = {0};
    uint32_t origin = 0;

    TEEC_SharedMemory pkSm;
    TEEC_SharedMemory smSm;

     /* ========================================================================
    [1] Connect to TEE
    ======================================================================== */
    result = TEEC_InitializeContext(NULL, &context);
    if (result != TEEC_SUCCESS)
    {
        printf("teec initial failed\n");
        goto cleanup_1;
    }

    /* ========================================================================
    [2] Open session with TEE application
    ======================================================================== */
    /* Open a Session with the TEE application. */
    /* MUST use TEEC_LOGIN_IDENTIFY method */
    operation.started = OPERATION_START_FLAG;
    operation.paramTypes = TEEC_PARAM_TYPES(TEEC_NONE, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    context.ta_path = k_qta_path;
    result = TEEC_OpenSession(&context, &session, &g_demoTemplateUuid, TEEC_LOGIN_IDENTIFY, NULL, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("teec open session failed\n");
        goto cleanup_2;
    }

     /* ========================================================================
    [3] Initialize the Shared Memory buffers
    ======================================================================== */
    /* Bulk input buffer. */
    pkSm.size = pkLen;
    pkSm.flags = TEEC_MEM_INPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    pkSm.buffer = (uint8_t *)pk;
    result = TEEC_RegisterSharedMemory(&context, &pkSm);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_3;
    }

    smSm.size = smLen;
    smSm.flags = TEEC_MEM_INPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    smSm.buffer = (uint8_t *)sm;
    result = TEEC_RegisterSharedMemory(&context, &smSm);
    if (result != TEEC_SUCCESS)
    {
        goto cleanup_4;
    }

    /* Start the dilithium verify operation within the TEE application. */
    operation.paramTypes =
        TEEC_PARAM_TYPES(TEEC_MEMREF_PARTIAL_INPUT, TEEC_MEMREF_PARTIAL_INPUT, TEEC_VALUE_OUTPUT, TEEC_NONE);

    operation.params[0].memref.parent = &pkSm;
    operation.params[0].memref.offset = 0;
    operation.params[0].memref.size = pkLen;

    operation.params[1].memref.parent = &smSm;
    operation.params[1].memref.offset = 0;
    operation.params[1].memref.size = smLen;
/* ========================================================================
    [4] Perform cryptographic operation dilithium verify
    ======================================================================== */
    result = TEEC_InvokeCommand(&session, CMD_DILITHIUM_VERIFY, &operation, &origin);
    *verify_result_ptr = operation.params[2].value.a;
    if (result != TEEC_SUCCESS) goto cleanup_5;


    TEEC_CloseSession(&session);
    return TEEC_SUCCESS;

cleanup_5:
    TEEC_ReleaseSharedMemory(&smSm);
cleanup_4:
    TEEC_ReleaseSharedMemory(&pkSm);
cleanup_3:
    TEEC_CloseSession(&session);
cleanup_2:
    TEEC_FinalizeContext(&context);
    return result;
cleanup_1:
    return result;
}

void bench_print(const char *label, double result)
{
    printf("BENCH: %s %*c = ", label, (int)(32 - strlen(label)), ' ');
    printf("%9.2f ms", result);

    if (result < 0)
    {
        printf(" (overflow or bad overhead estimation)\n");
    }
    else
    {
        printf("\n");
    }
}

TEEC_Result dilithium_testbench()
{
    TEEC_Context context = {0};
    TEEC_Session session = {0};
    TEEC_Result result;
    TEEC_Operation operation = {0};
    uint32_t origin = 0;

    TEEC_SharedMemory timeCostSM;
    uint32_t timeCost[13] = {0};

     /* ========================================================================
    [1] Connect to TEE
    ======================================================================== */
    result = TEEC_InitializeContext(NULL, &context);
    if (result != TEEC_SUCCESS)
    {
        printf("dilithium_testbench: teec initial failed\n");
        goto cleanup_1;
    }

    /* ========================================================================
    [2] Open session with TEE application
    ======================================================================== */
    /* Open a Session with the TEE application. */
    /* MUST use TEEC_LOGIN_IDENTIFY method */
    operation.started = OPERATION_START_FLAG;
    operation.paramTypes = 
        TEEC_PARAM_TYPES(TEEC_NONE, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    context.ta_path = k_qta_path;
    result = TEEC_OpenSession(&context, &session, &g_demoTemplateUuid, TEEC_LOGIN_IDENTIFY, NULL, &operation, &origin);
    if (result != TEEC_SUCCESS)
    {
        printf("dilithium_testbench: teec open session failed\n");
        goto cleanup_2;
    }

     /* ========================================================================
    [3] Initialize the Shared Memory buffers
    ======================================================================== */
    /* Bulk input buffer. */
    timeCostSM.size = sizeof(timeCost);
    timeCostSM.flags = TEEC_MEM_OUTPUT;
    /* Use TEE Client API to register the underlying memory buffer. */
    timeCostSM.buffer = (uint8_t *)timeCost;
    result = TEEC_RegisterSharedMemory(&context, &timeCostSM);
    if (result != TEEC_SUCCESS)
    {
        printf("dilithium_testbench: teec register shared memory failed\n");
        goto cleanup_3;
    }

    /* Start the dilithium verify operation within the TEE application. */
    operation.paramTypes =
        TEEC_PARAM_TYPES(TEEC_MEMREF_PARTIAL_OUTPUT, TEEC_NONE, TEEC_NONE, TEEC_NONE);

    operation.params[0].memref.parent = &timeCostSM;
    operation.params[0].memref.offset = 0;
    operation.params[0].memref.size = sizeof(timeCost);

/* ========================================================================
    [4] Perform cryptographic operation dilithium testbench
    ======================================================================== */
    clock_t before = clock();
    result = TEEC_InvokeCommand(&session, CMD_DILITHIUM_TESTBENCH, &operation, &origin);
    clock_t after = clock();
    // microsecond
    uint32_t t_res = (after - before) * 1000000 / CLOCKS_PER_SEC;
    printf("REE cost: %u us\n", t_res);

    printf("Bench round: %u.\n", timeCost[0]);
    printf("+---    Bench module  --->  Dilithium Level %d.\n", 2);
    bench_print("crypto_sign_keypair", timeCost[1]);
    bench_print("crypto_sign", timeCost[2]);
    bench_print("crypto_sign_open", timeCost[3]);
    printf("+---    Bench module  --->  Dilithium Level %d.\n", 3);
    bench_print("crypto_sign_keypair", timeCost[4]);
    bench_print("crypto_sign", timeCost[5]);
    bench_print("crypto_sign_open", timeCost[6]);
    printf("+---    Bench module  --->  Dilithium Level %d.\n", 5);
    bench_print("crypto_sign_keypair", timeCost[7]);
    bench_print("crypto_sign", timeCost[8]);
    bench_print("crypto_sign_open", timeCost[9]);
    printf("+---    Bench module  --->  Rng.\n");
    bench_print("TEE_GenerateRandom", timeCost[10]);
    bench_print("shake128", timeCost[11]);
    bench_print("shake256", timeCost[12]);
    
    
    if (result != TEEC_SUCCESS)
    {
        printf("dilithium testbench failed\n");
        goto cleanup_4;
    }
    else
    {
        printf("dilithium testbench succeeded\n");
    }


    TEEC_CloseSession(&session);
    return TEEC_SUCCESS;

cleanup_4:
    TEEC_ReleaseSharedMemory(&timeCostSM);
cleanup_3:
    TEEC_CloseSession(&session);
cleanup_2:
    TEEC_FinalizeContext(&context);
    return result;
cleanup_1:
    return result;
}

int main(void)
{
    get_ta_version();

    dilithium_testbench();
    return 0;

}
